#pragma once

#include "filters/simple_filter_engine.h"
#include "core/packet_dispatcher.h"
#include "core/packet.h"

#include <atomic>
#include <thread>
#include <functional>

class FilterWorker {
 public:
    using PacketCallBack = std::function<void(Packet&&, const FilterResult&)>;

    FilterWorker(size_t workerId, std::shared_ptr<IFilterEngine> filterEngine, PacketDispatcher::PacketQueue& inputQueue);
    ~FilterWorker();

    /// Disable copying and assignment
    FilterWorker(const FilterWorker&) = delete;
    FilterWorker& operator=(const FilterWorker&) = delete;

    /// Start and stop work thread
    bool start();
    void stop();
    bool isRunning() const { return m_running.load(); }

    /// Get worker thread id
    size_t getId() const { return m_workerId; }

    void setBatchSize(size_t batchSize) { m_batchSize = batchSize; }

    /// Used to forward received packets
    void setOutputCallback(PacketCallBack callback) { m_outputCallback = callback; }

    /// Statistics
    struct WorkerStats {
        size_t worker_id;
        uint64_t processed_packets = 0;
        uint64_t accepted_packets = 0;
        uint64_t dropped_packets = 0;
        uint64_t queue_empty_cycles = 0;
        uint64_t batch_processed = 0;
        double average_processing_time_us = 0.0;
        size_t current_queue_size = 0;
    };

    WorkerStats getStats() const;

 private:
    void workerThread();
    void processSinglePacket();
    void processBatch();
    void handleEmptyQueue();

    const size_t m_workerId;
    std::shared_ptr<IFilterEngine> m_filterEngine;
    PacketDispatcher::PacketQueue& m_inputQueue;
    std::thread m_workerThread;
    std::atomic<bool> m_running{false};
    std::atomic<bool> m_shouldStop{false};

    PacketCallBack m_outputCallback;
    size_t m_batchSize = 16;

    /// Statistics
    mutable std::mutex m_statsMutex;
    WorkerStats m_stats;
    std::atomic<uint64_t> m_totalProcessingTime{0};  /// ns
};

class FilterThreadPool {
 public:
    FilterThreadPool(size_t threadCount, std::shared_ptr<IFilterEngine> filterEngine, std::shared_ptr<PacketDispatcher> dispatcher);
    ~FilterThreadPool();

    /// Disable copying and assignment
    FilterThreadPool(const FilterThreadPool&) = delete;
    FilterThreadPool& operator=(const FilterThreadPool&) = delete;

    /// @brief Start all worker threads
    bool start();
    /// @brief Stop all worker threads
    void stop();

    bool isRunning() const { return m_running.load(std::memory_order::acquire); }

    /// Distribute tasks to different workers using round-robin or hash-based scheduling
    void setOutputCallback(FilterWorker::PacketCallBack callback);

    /// Statistics
    std::vector<FilterWorker::WorkerStats> getAllStats() const;
    size_t getActiveThreadCount() const;

    struct PoolStats {
        size_t total_threads;
        size_t active_threads;
        uint64_t total_processed_packets;
        uint64_t total_accepted_packets;
        uint64_t total_dropped_packets;
        double overall_throughput_pps; /// packet per second
    };

    PoolStats getPoolStats() const;

 private:
    void createWorkers();

    std::vector<std::unique_ptr<FilterWorker>> m_workers;
    std::shared_ptr<IFilterEngine> m_filterEngine;
    std::shared_ptr<PacketDispatcher> m_dispatcher;
    size_t m_threadCount;
    std::atomic<bool> m_running{false};
};